use base64::prelude::BASE64_STANDARD;
use base64::Engine;
use std::mem;

include!("consts.rs");

pub struct ZUCKeyWordIter {
    s: [u32; 16],
    r: [u32; 2],
    x: [u32; 4],
}

impl ZUCKeyWordIter {
    fn mod_add(a: u32, b: u32) -> u32 {
        let c = a.wrapping_add(b);
        (c & 0x7FFFFFFF).wrapping_add(c >> 31)
    }

    fn mod_mul(a: u32, b: u32) -> u32 {
        // TODO impl a better mul
        (((a as u64) * (b as u64)) % 2147483647u64) as u32
    }

    fn l1(bits: u32) -> u32 {
        bits ^ bits.rotate_left(2)
            ^ bits.rotate_left(10)
            ^ bits.rotate_left(18)
            ^ bits.rotate_left(24)
    }

    fn l2(bits: u32) -> u32 {
        bits ^ bits.rotate_left(8)
            ^ bits.rotate_left(14)
            ^ bits.rotate_left(22)
            ^ bits.rotate_left(30)
    }

    fn s_box_trans(u: u32) -> u32 {
        const BIT_OFFSET: [u8; 4] = [24, 16, 8, 0];
        BIT_OFFSET.iter().enumerate().fold(0, |res, (i, &offset)| {
            res | ((S[i % 2][(u >> offset & 0xFF) as usize] as u32) << offset)
        })
    }

    fn bit_reconstruction(&mut self) {
        const HIGH_MAKE: u32 = 0x7fff8000;
        const LOW_MAKE: u32 = 0x0000ffff;
        self.x[0] = ((self.s[15] & HIGH_MAKE) << 1) | (self.s[14] & LOW_MAKE);
        self.x[1] = ((self.s[11] & LOW_MAKE) << 16) | (self.s[9] >> 15);
        self.x[2] = ((self.s[7] & LOW_MAKE) << 16) | (self.s[5] >> 15);
        self.x[3] = ((self.s[2] & LOW_MAKE) << 16) | (self.s[0] >> 15);
    }
    fn f(&mut self) -> u32 {
        let w = (self.x[0] ^ self.r[0]).wrapping_add(self.r[1]);
        let w1 = self.r[0].wrapping_add(self.x[1]);
        let w2 = self.r[1] ^ self.x[2];
        self.r[0] = Self::s_box_trans(Self::l1((w1 << 16) | (w2 >> 16)));
        self.r[1] = Self::s_box_trans(Self::l2((w2 << 16) | (w1 >> 16)));
        w
    }
    fn lfsr(&mut self, u: u32) {
        const FACTOR: [u32; 5] = [1 << 15, 1 << 17, 1 << 21, 1 << 20, (1 << 8) + 1];
        const S_IDX: [usize; 5] = [15, 13, 10, 4, 0];
        let v = S_IDX
            .iter()
            .zip(FACTOR.iter())
            .fold(0, |sum, (&idx, &factor)| {
                Self::mod_add(Self::mod_mul(self.s[idx], factor), sum)
            });
        for i in 0..15 {
            self.s[i] = self.s[i + 1];
        }
        self.s[15] = Self::mod_add(v, u);
        if self.s[15] == 0 {
            self.s[15] = 0x7fffffff;
        }
    }

    pub fn new(key: &[u8], iv: &[u8]) -> Self {
        let mut s = [0u32; 16];
        key.iter()
            .zip(iv)
            .zip(D)
            .enumerate()
            .for_each(|(i, ((&key, &iv), d))| s[i] = (key as u32) << 23 | d << 8 | (iv as u32));
        let mut zuc = Self {
            s,
            r: [0; 2],
            x: [0; 4],
        };
        (0..32).for_each(|_| {
            zuc.bit_reconstruction();
            let w = zuc.f();
            zuc.lfsr(w >> 1);
        });
        let _ = zuc.next();
        zuc
    }
}

impl Iterator for ZUCKeyWordIter {
    type Item = u32;
    fn next(&mut self) -> Option<Self::Item> {
        self.bit_reconstruction();
        let z = self.f() ^ self.x[3];
        self.lfsr(0);
        Some(z)
    }
}

fn pkcs7_padding(
    input: impl Iterator<Item = u8>,
    block_size: usize,
    origin_length: usize,
) -> impl Iterator<Item = u8> {
    let padding_size = block_size - (origin_length % block_size);
    input.chain(std::iter::repeat(padding_size as u8).take(padding_size))
}

fn bytes_to_32words(input: impl Iterator<Item = u8>) -> impl Iterator<Item = u32> {
    input
        .scan(([0; BLOCK_SIZE], 0), |(buf, buf_size), x| {
            buf[*buf_size] = x;
            *buf_size += 1;
            if *buf_size >= BLOCK_SIZE {
                *buf_size = 0;
                let word = u32::from_be_bytes(mem::replace(buf, [0; BLOCK_SIZE]));
                return Some(Some(word));
            }
            Some(None)
        })
        .filter_map(|x| x)
}
const BLOCK_SIZE: usize = 4;

pub fn encryption(input: String) -> String {
    let key = [0u8; 16];
    let iv = [1u8; 16];
    let input = input.as_bytes();
    let input = pkcs7_padding(input.iter().copied(), BLOCK_SIZE, input.len());
    let input = bytes_to_32words(input);
    let dst = input
        .zip(ZUCKeyWordIter::new(&key, &iv))
        .map(|(ibs, key)| (ibs ^ key).to_be_bytes())
        .flatten()
        .collect::<Vec<_>>();
    BASE64_STANDARD.encode(dst.as_slice())
}
